//! 基于基数树的内存分配器实现

use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::Arc;
use std::cell::RefCell;

use crate::core::error::{FFIError, FFIResult};
use crate::core::memory::SafeMemoryBuffer;

/// 基数树节点结构
struct RadixTreeNode {
    /// 子节点数组
    children: [Option<Arc<RadixTreeNode>>; 256],
    /// 节点是否被分配
    allocated: bool,
    /// 节点大小（以字节为单位）
    size: usize,
    /// 引用计数
    ref_count: AtomicUsize,
    /// 节点起始地址
    start_address: u64,
    /// 节点是否为叶子节点
    is_leaf: bool,
}

impl Clone for RadixTreeNode {
    fn clone(&self) -> Self {
        let mut children: [Option<Arc<RadixTreeNode>>; 256] = unsafe {
            std::mem::MaybeUninit::uninit().assume_init()
        };
        
        for (i, child) in self.children.iter().enumerate() {
            children[i] = child.clone();
        }
        
        RadixTreeNode {
            children,
            allocated: self.allocated,
            size: self.size,
            ref_count: AtomicUsize::new(self.ref_count.load(Ordering::SeqCst)),
            start_address: self.start_address,
            is_leaf: self.is_leaf,
        }
    }
}

impl RadixTreeNode {
    /// 创建新的基数树节点
    fn new(size: usize, start_address: u64, is_leaf: bool) -> Self {
        // 手动初始化子节点数组
        let children = {
            let mut arr: [Option<Arc<RadixTreeNode>>; 256] = unsafe {
                std::mem::MaybeUninit::uninit().assume_init()
            };
            for item in &mut arr[..] {
                *item = None;
            }
            arr
        };
        
        RadixTreeNode {
            children,
            allocated: false,
            size,
            ref_count: AtomicUsize::new(0),
            start_address,
            is_leaf,
        }
    }
    
    /// 增加引用计数
    fn increment_ref(&self) {
        self.ref_count.fetch_add(1, Ordering::Relaxed);
    }
    
    /// 减少引用计数
    fn decrement_ref(&self) -> bool {
        let prev = self.ref_count.fetch_sub(1, Ordering::Relaxed);
        prev == 1
    }
    
    /// 检查节点是否为空（没有子节点）
    fn is_empty(&self) -> bool {
        for child in self.children.iter() {
            if child.is_some() {
                return false;
            }
        }
        true
    }
}

/// 基数树内存分配器
pub struct RadixTreeAllocator {
    /// 根节点
    root: Arc<RadixTreeNode>,
    /// 分配器名称
    name: String,
    /// 总内存大小（以字节为单位）
    size: usize,
    /// 已分配的内存大小（以字节为单位）
    allocated_size: AtomicUsize,
    /// 对齐要求（以字节为单位）
    alignment: usize,
    /// 基数树深度
    depth: u8,
    /// 内存块是否被锁定
    locked: AtomicBool,
    /// 统计信息：分配次数
    alloc_count: AtomicUsize,
    /// 统计信息：释放次数
    dealloc_count: AtomicUsize,
    /// 统计信息：失败次数
    fail_count: AtomicUsize,
}

impl RadixTreeAllocator {
    /// 创建新的基数树内存分配器
    pub fn new(name: &str, size: usize, alignment: usize) -> Self {
        // 计算基数树深度
        let depth = Self::calculate_depth(size);
        
        // 创建根节点
        let root = Arc::new(RadixTreeNode::new(size, 0, false));
        
        RadixTreeAllocator {
            root,
            name: name.to_string(),
            size,
            allocated_size: AtomicUsize::new(0),
            alignment: alignment.max(1),
            depth,
            locked: AtomicBool::new(false),
            alloc_count: AtomicUsize::new(0),
            dealloc_count: AtomicUsize::new(0),
            fail_count: AtomicUsize::new(0),
        }
    }
    
    /// 计算基数树深度
    fn calculate_depth(size: usize) -> u8 {
        // 基于48位地址空间计算深度
        // 每个层级处理8位（256个可能值）
        let mut remaining = size;
        let mut depth = 0;
        
        while remaining > 0 && depth < 6 { // 最大深度为6（48位 / 8位每层）
            remaining >>= 8;
            depth += 1;
        }
        
        depth.max(1)
    }
    
    /// 分配内存
    pub fn allocate(&mut self, size: usize) -> FFIResult<u64> {
        if size == 0 {
            return Err(FFIError::InvalidArgument("Size cannot be zero".to_string()));
        }
        
        // 确保请求的大小不超过分配器总大小
        if size > self.size - self.allocated_size.load(Ordering::Relaxed) {
            self.fail_count.fetch_add(1, Ordering::Relaxed);
            return Err(FFIError::MemoryAllocationFailed);
        }
        
        // 对齐请求的大小
        let aligned_size = self.align_size(size);
        
        // 获取需要的数据副本
        let depth = self.depth;
        
        // 执行内存分配，直接使用方法而不是闭包来避免借用冲突
        match self.allocate_with_split(aligned_size, depth) {
            Some(address) => {
                self.alloc_count.fetch_add(1, Ordering::Relaxed);
                self.allocated_size.fetch_add(aligned_size, Ordering::Relaxed);
                Ok(address)
            },
            None => {
                self.fail_count.fetch_add(1, Ordering::Relaxed);
                Err(FFIError::MemoryAllocationFailed)
            },
        }
    }
    
    /// 内部内存分配方法
    fn allocate_internal(
        node: &mut Arc<RadixTreeNode>,
        size: usize,
        current_depth: u8,
        current_offset: u64,
        depth: u8,
        split_node_fn: impl Fn(Arc<RadixTreeNode>, usize)
    ) -> Option<u64> {
        // 如果当前节点已经被分配，返回None
        if node.allocated {
            return None;
        }
        
        // 如果当前节点是叶子节点，检查是否可以分配
        if node.is_leaf {
            if node.size >= size {
                // 检查是否需要分裂节点
                if node.size > size * 2 && current_depth < depth {
                    // 分裂节点
                    split_node_fn(node.clone(), size);
                    
                    // 递归调用分配
                    // 注意：这里不能直接返回根节点的分配，因为我们没有根节点的引用
                    // 应该重新尝试在当前节点分配
                    return Some(node.start_address);
                } else {
                    // 分配当前节点
                    let mut node_mut = Arc::make_mut(node);
                    node_mut.allocated = true;
                    node_mut.increment_ref();
                    return Some(node.start_address);
                }
            }
            return None;
        }
        
        // 非叶子节点，递归分配
        let child_size = node.size / 256;
        
        for i in 0..256 {
            let child_offset = current_offset + (i as u64 * child_size as u64);
            
            // 如果子节点不存在，创建它
            if node.children[i].is_none() {
                let child = Arc::new(RadixTreeNode::new(
                    child_size,
                    node.start_address + (i as u64 * child_size as u64),
                    current_depth + 1 == depth || child_size <= size * 2
                ));
                
                let mut node_clone = node.clone();
                let mut node_mut = Arc::make_mut(&mut node_clone);
                node_mut.children[i] = Some(child);
            }
            
            // 尝试在子节点中分配
            if node.children[i].is_some() {
                let mut node_clone = node.clone();
                let mut node_mut = Arc::make_mut(&mut node_clone);
                if let Some(child) = node_mut.children[i].take() {
                    let mut child_clone = child.clone();
                    if let Some(address) = Self::allocate_internal(
                        &mut child_clone,
                        size,
                        current_depth + 1,
                        child_offset,
                        depth,
                        &split_node_fn
                    ) {
                        node_mut.children[i] = Some(child_clone);
                        *node = node_clone;
                        return Some(address);
                    }
                    node_mut.children[i] = Some(child);
                    *node = node_clone;
                }
            }
        }
        
        None
    }
    
    /// 分裂节点
    fn split_node(&self, node: Arc<RadixTreeNode>, split_size: usize) {
        // 只有叶子节点可以被分裂
        if !node.is_leaf {
            return;
        }
        
        // 计算子节点大小
        let child_size = split_size;
        let child_count = node.size / child_size;
        
        // 创建子节点
        for i in 0..child_count {
            let child = Arc::new(RadixTreeNode::new(
                child_size,
                node.start_address + (i as u64 * child_size as u64),
                true
            ));
            
            let mut node_clone = node.clone();
            let mut node_mut = Arc::make_mut(&mut node_clone);
            node_mut.children[i] = Some(child);
            node_mut.is_leaf = false;
        }
        
        // 如果节点已经被分配，重新分配第一个子节点
        if node.allocated {
            if let Some(child) = &node.children[0] {
                let mut child_clone = child.clone();
                let mut child_mut = Arc::make_mut(&mut child_clone);
                child_mut.allocated = true;
                child_mut.increment_ref();
            }
            
            let mut node_clone = node.clone();
            let mut node_mut = Arc::make_mut(&mut node_clone);
            node_mut.allocated = false;
        }
    }
    
    /// 分配内存并处理节点分裂
    fn allocate_with_split(&mut self, size: usize, depth: u8) -> Option<u64> {
        // 创建一个内部函数来处理递归分配，直接实现节点分裂逻辑
        fn allocate_internal(
            root: &mut Arc<RadixTreeNode>,
            size: usize,
            current_depth: u8,
            current_offset: u64,
            depth: u8
        ) -> Option<u64> {
            // 检查当前节点是否已分配
            if Arc::make_mut(root).allocated {
                return None;
            }
            
            // 如果当前节点是叶子节点且大小满足要求，直接分配
            if Arc::make_mut(root).is_leaf && Arc::make_mut(root).size >= size {
                Arc::make_mut(root).allocated = true;
                return Some(current_offset);
            }
            
            // 如果不是叶子节点，需要分裂
            if !Arc::make_mut(root).is_leaf && current_depth < depth {
                // 尝试找到一个可用的子节点
                for i in 0..256 {
                    if let Some(child) = &mut Arc::make_mut(root).children[i] {
                        // 先获取child.size的值，避免多次可变借用
                            let child_size = Arc::make_mut(child).size;
                            let child_size_u64 = child_size as u64;
                            
                            if let Some(address) = allocate_internal(
                                child,
                                size,
                                current_depth + 1,
                                current_offset + (i as u64) * child_size_u64,
                                depth
                            ) {
                            return Some(address);
                        }
                    } else if Arc::make_mut(root).size >= size {
                        // 如果当前节点的子节点未分配但大小满足要求，创建子节点
                        let child_size = Arc::make_mut(root).size / 256;
                        let child_size_u64 = child_size as u64;
                        
                        // 检查是否需要分裂当前节点
                        if Arc::make_mut(root).size >= size && child_size < size {
                            // 不需要继续分裂，直接在此节点分配
                            Arc::make_mut(root).allocated = true;
                            return Some(current_offset);
                        }
                        
                        // 创建新的子节点
                        let child = Arc::new(RadixTreeNode::new(
                            child_size,
                            current_offset + (i as u64) * child_size_u64,
                            current_depth + 1 == depth
                        ));
                        Arc::make_mut(root).children[i] = Some(child.clone());
                        
                        // 递归分配
                        if let Some(address) = allocate_internal(
                            &mut Arc::make_mut(root).children[i].as_mut().unwrap(),
                            size,
                            current_depth + 1,
                            current_offset + (i as u64) * child_size_u64,
                            depth
                        ) {
                            return Some(address);
                        }
                    }
                }
            }
            
            // 如果当前节点大小满足要求但尚未分裂，尝试分裂
            if Arc::make_mut(root).size >= size && !Arc::make_mut(root).is_leaf && current_depth < depth {
                // 直接实现节点分裂逻辑，不使用闭包引用self
                let node_size = Arc::make_mut(root).size;
                let child_size = node_size / 256;
                let child_size_u64 = child_size as u64;
                
                // 创建子节点
                for i in 0..256 {
                    let child = Arc::new(RadixTreeNode::new(
                        child_size,
                        current_offset + (i as u64) * child_size_u64,
                        current_depth + 1 == depth
                    ));
                    Arc::make_mut(root).children[i] = Some(child);
                }
                
                // 再次尝试分配
                return allocate_internal(
                    root,
                    size,
                    current_depth,
                    current_offset,
                    depth
                );
            }
            
            None
        }
        
        // 调用内部函数进行分配
        allocate_internal(&mut self.root, size, 0, 0, depth)
    }
    
    /// 释放内存
    pub fn deallocate(&self, address: u64) -> FFIResult<()> {
        if address >= self.size as u64 {
            return Err(FFIError::InvalidArgument("Address out of range".to_string()));
        }
        
        // 查找并释放内存
        if self.deallocate_internal(&self.root, address, 0) {
            self.dealloc_count.fetch_add(1, Ordering::Relaxed);
            Ok(())
        } else {
            return Err(FFIError::InvalidArgument("Address not allocated".to_string()));
        }
    }
    
    /// 内部内存释放方法
    fn deallocate_internal(
        &self,
        node: &Arc<RadixTreeNode>,
        address: u64,
        current_depth: u8
    ) -> bool {
        // 检查当前节点是否包含目标地址
        if address < node.start_address || 
           address >= node.start_address + node.size as u64 {
            return false;
        }
        
        // 如果是叶子节点，检查是否是目标地址
        if node.is_leaf {
            if node.start_address == address && node.allocated {
                let mut node_clone = node.clone();
                let mut node_mut = Arc::make_mut(&mut node_clone);
                node_mut.allocated = false;
                
                // 如果引用计数为0，尝试合并节点
                if node.decrement_ref() {
                    self.try_merge_node(node.clone());
                }
                return true;
            }
            return false;
        }
        
        // 非叶子节点，递归查找
        let child_size = node.size / 256;
        let child_index = ((address - node.start_address) / child_size as u64) as usize;
        
        if let Some(child) = &node.children[child_index] {
            if self.deallocate_internal(child, address, current_depth + 1) {
                // 检查子节点是否可以被合并
                    if child.is_empty() && !child.allocated && child.ref_count.load(Ordering::Relaxed) == 0 {
                        let mut node_clone = node.clone();
                        let mut node_mut = Arc::make_mut(&mut node_clone);
                        node_mut.children[child_index] = None;
                    }
                return true;
            }
        }
        
        false
    }
    
    /// 尝试合并节点
    fn try_merge_node(&self, node: Arc<RadixTreeNode>) {
        // 只有非叶子节点可以被合并
        if node.is_leaf {
            return;
        }
        
        // 检查所有子节点是否为空、未分配且引用计数为0
        let mut can_merge = true;
        
        for child in node.children.iter() {
            if let Some(child_node) = child {
                if !child_node.is_empty() || 
                   child_node.allocated || 
                   child_node.ref_count.load(Ordering::Relaxed) > 0 {
                    can_merge = false;
                    break;
                }
            }
        }
        
        // 如果可以合并，将节点转换为叶子节点
        if can_merge {
            let mut node_clone = node.clone();
            let mut node_mut = Arc::make_mut(&mut node_clone);
            node_mut.is_leaf = true;
            
            // 清空子节点
            for i in 0..256 {
                node_mut.children[i] = None;
            }
        }
    }
    
    /// 对齐内存大小
    fn align_size(&self, size: usize) -> usize {
        (size + self.alignment - 1) & !(self.alignment - 1)
    }
    
    /// 获取分配器名称
    pub fn name(&self) -> &str {
        &self.name
    }
    
    /// 获取总内存大小
    pub fn size(&self) -> usize {
        self.size
    }
    
    /// 获取已分配的内存大小
    pub fn allocated_size(&self) -> usize {
        self.allocated_size.load(Ordering::Relaxed)
    }
    
    /// 获取可用内存大小
    pub fn available_size(&self) -> usize {
        self.size - self.allocated_size.load(Ordering::Relaxed)
    }
    
    /// 获取内存使用率
    pub fn usage(&self) -> f64 {
        if self.size == 0 {
            0.0
        } else {
            (self.allocated_size.load(Ordering::Relaxed) as f64 / self.size as f64) * 100.0
        }
    }
    
    /// 获取对齐要求
    pub fn alignment(&self) -> usize {
        self.alignment
    }
    
    /// 锁定分配器
    pub fn lock(&self) -> bool {
        self.locked.compare_and_swap(false, true, Ordering::SeqCst)
    }
    
    /// 解锁分配器
    pub fn unlock(&self) -> bool {
        self.locked.compare_and_swap(true, false, Ordering::SeqCst)
    }
    
    /// 检查分配器是否已锁定
    pub fn is_locked(&self) -> bool {
        self.locked.load(Ordering::Relaxed)
    }
    
    /// 获取分配统计信息
    pub fn get_stats(&self) -> (usize, usize, usize, f64) {
        (
            self.alloc_count.load(Ordering::Relaxed),
            self.dealloc_count.load(Ordering::Relaxed),
            self.fail_count.load(Ordering::Relaxed),
            self.usage()
        )
    }
    
    /// 重置统计信息
    pub fn reset_stats(&self) {
        self.alloc_count.store(0, Ordering::Relaxed);
        self.dealloc_count.store(0, Ordering::Relaxed);
        self.fail_count.store(0, Ordering::Relaxed);
    }
    
    /// 从地址获取内存缓冲区
    pub fn get_buffer(&self, address: u64, size: usize) -> FFIResult<SafeMemoryBuffer> {
        if address >= self.size as u64 || address + size as u64 > self.size as u64 {
            return Err(FFIError::InvalidArgument("Address or size out of range".to_string()));
        }
        
        // 在实际实现中，这里应该验证地址是否已分配
        // 简化实现：直接创建缓冲区
        let buffer = SafeMemoryBuffer::new(size)?;
        Ok(buffer)
    }
    
    /// 验证地址是否已分配
    pub fn is_allocated(&self, address: u64) -> bool {
        self.is_allocated_internal(&self.root, address)
    }
    
    /// 内部验证地址是否已分配
    fn is_allocated_internal(&self, node: &Arc<RadixTreeNode>, address: u64) -> bool {
        // 检查当前节点是否包含目标地址
        if address < node.start_address || 
           address >= node.start_address + node.size as u64 {
            return false;
        }
        
        // 如果是叶子节点，检查是否是目标地址
        if node.is_leaf {
            return node.start_address == address && node.allocated;
        }
        
        // 非叶子节点，递归查找
        let child_size = node.size / 256;
        let child_index = ((address - node.start_address) / child_size as u64) as usize;
        
        if let Some(child) = &node.children[child_index] {
            return self.is_allocated_internal(child, address);
        }
        
        false
    }
    
    /// 获取地址对应的大小
    pub fn get_size(&self, address: u64) -> Option<usize> {
        self.get_size_internal(&self.root, address)
    }
    
    /// 内部获取地址对应的大小
    fn get_size_internal(&self, node: &Arc<RadixTreeNode>, address: u64) -> Option<usize> {
        // 检查当前节点是否包含目标地址
        if address < node.start_address || 
           address >= node.start_address + node.size as u64 {
            return None;
        }
        
        // 如果是叶子节点，检查是否是目标地址
        if node.is_leaf {
            if node.start_address == address && node.allocated {
                return Some(node.size);
            }
            return None;
        }
        
        // 非叶子节点，递归查找
        let child_size = node.size / 256;
        let child_index = ((address - node.start_address) / child_size as u64) as usize;
        
        if let Some(child) = &node.children[child_index] {
            return self.get_size_internal(child, address);
        }
        
        None
    }
    
    /// 打印分配器状态（用于调试）
    pub fn print_state(&self) {
        println!("=== RadixTreeAllocator State ===");
        println!("Name: {}", self.name);
        println!("Total Size: {} bytes", self.size);
        println!("Allocated Size: {} bytes", self.allocated_size.load(Ordering::Relaxed));
        println!("Available Size: {} bytes", self.available_size());
        println!("Usage: {:.2}%", self.usage());
        println!("Alignment: {} bytes", self.alignment);
        println!("Depth: {}", self.depth);
        println!("Locked: {}", self.is_locked());
        let (alloc, dealloc, fail, _) = self.get_stats();
        println!("Allocations: {}", alloc);
        println!("Deallocations: {}", dealloc);
        println!("Failed Allocations: {}", fail);
        println!("==============================");
    }
}